
import queue
import time
from abc import ABC
from collections.abc import Generator
from typing import Any, Optional, Union
from uuid import UUID
from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import LLMResult
from openai import AuthenticationError, RateLimitError, OpenAIError, APIStatusError, InternalServerError
from llmchat.model import ApiResponse
from llmchat.model.ApiResponse import ErrorCode
from loguru import logger


class ChainStreamHandler(BaseCallbackHandler, ABC):
    """流式回调处理器"""
    end = "[DONE]"
    end_rsp = ApiResponse.success(end)

    def __init__(self, method_name: str = "", requestId: str = ""):
        self.tokens = queue.Queue()
        self.start_time = time.time()
        self.method_name = method_name
        self.requestId = requestId
        super(BaseCallbackHandler, self).__init__()

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        token = token.replace("```json\n", "").replace("\n```", "")
        logger.error("sse处理，接收到数据" + self.requestId + "):" + token)
        self.tokens.put(token)

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        end_time = time.time()
        result = "响应内容异常"
        if len(response.generations) > 0:
            if len(response.generations[0]) > 0:
                result = response.generations[0][0].text
                result = result.replace(" ", "").replace("\n", "")

        logger.info(self.method_name + "sse运行时间(" + self.requestId + ")：%.2f秒" % (end_time - self.start_time))
        logger.info(self.method_name + "sse应答内容(" + self.requestId + "):" + result)

    def on_llm_error(
            self,
            error: Union[Exception, KeyboardInterrupt],
            *,
            run_id: UUID,
            parent_run_id: Optional[UUID] = None,
            **kwargs: Any,
    ) -> None:
        logger.error("sse处理异常(" + self.requestId + "):" + str(error))
        self.tokens.put(error)
        self.tokens.put(self.end)

    def on_chain_error(
            self,
            error: Union[Exception, KeyboardInterrupt],
            *,
            run_id: UUID,
            parent_run_id: Optional[UUID] = None,
            **kwargs: Any,
    ) -> None:
        logger.error("sse处理异常(" + self.requestId + "):" + str(error))
        self.tokens.put(error)
        self.tokens.put(self.end)

    def on_tool_error(
            self,
            error: Union[Exception, KeyboardInterrupt],
            *,
            run_id: UUID,
            parent_run_id: Optional[UUID] = None,
            **kwargs: Any,
    ) -> None:
        logger.error("sse处理异常(" + self.requestId + "):" + str(error))
        self.tokens.put(error)
        self.tokens.put(self.end)

    def generate_tokens(self) -> Generator:
        while True:
            tk = self.tokens.get()
            try:
                if tk == self.end:
                    yield f"data: {self.end_rsp}\n\n"
                    return
                elif isinstance(tk, APIStatusError):
                    error_enum = ErrorCode.AI_DEFAULT_EXCEPTION
                    if isinstance(tk, AuthenticationError):
                        error_enum = ErrorCode.AI_ACCOUNT_ERROR
                    elif isinstance(tk, RateLimitError):
                        error_enum = ErrorCode.AI_RATE_LIMIT_ERROR
                    error_msg = error_enum.message + ",status_code="+str(tk.status_code)+",message="+tk.message
                    logger.error("sse响应异常(" + self.requestId + ")"+error_msg)
                    error_rsp = ApiResponse.error(error_enum.code, error_msg)
                    yield f"data: {error_rsp}\n\n"
                else:
                    data = ApiResponse.success(tk)
                    yield f"data: {data}\n\n"
            except Exception as ex:
                error_msg = str(ex)
                logger.error("sse处理异常(" + self.requestId + "):" + error_msg)
                error = ApiResponse.error(code=ErrorCode.AI_DEFAULT_EXCEPTION.code, msg=error_msg)
                yield f"data: {error}\n\n"
                yield f"data: {self.end_rsp}\n\n"
                return
